Chuyển đổi từ Chế độ Eager của PyTorch sang Triton yêu cầu thay đổi cách nhìn từ tensor như một đối tượng nguyên vẹn sang xem chúng như tập hợp các khối riêng biệt, dễ quản lý khối hoặc ô.

1. Tensor của PyTorch so với Triton

Điều quan trọng là phải phân biệt tensor của Triton với tensor của PyTorch. Một tensor của PyTorch là một đối tượng Python ở phía máy chủ bao bọc thông tin hình dạng, kiểu dữ liệu, thiết bị, bước nhảy và thông tin lưu trữ. Ngược lại, Triton làm việc với con trỏ dữ liệu thô trong các khối bộ nhớ cụ thể, cho phép tối ưu hóa ở mức thấp hơn nhiều.

2. Cái nút thắt Eager

Trong thực thi ngay lập tức tiêu chuẩn, mỗi thao tác (ví dụ: cộng rồi sau đó áp dụng ReLU) đều cần khởi chạy kernel riêng biệt và một lượt đi về bộ nhớ toàn cục. Đây là điểm nghẽn chính trong tính toán GPU hiện đại. Triton vượt qua điều này bằng cách gộp các thao tác trong một kernel duy nhất xử lý các khối dữ liệu (ví dụ: 128, 256 hoặc 512 phần tử) trực tiếp trong bộ nhớ trên chip.

3. Mô hình dựa trên khối

Thay vì tư duy cấp số vô hướng như các luồng CUDA, Triton sử dụng SPMD (Chương trình đơn, Dữ liệu đa) ở cấp độ khối. Bạn viết một kernel, và Triton khởi chạy nhiều phiên bản trên một lưới. Mỗi phiên bản sử dụng program_id để tính toán khối bộ nhớ nào mà nó sở hữu.

4. Thiết lập môi trường

Để bắt đầu, cài đặt Triton trong môi trường sạch (dùng Conda hoặc venv) để đảm bảo không có xung đột phụ thuộc với các công cụ CUDA hiện có: pip install triton.